package com.share.ai.langchain4j.web;


import com.share.ai.langchain4j.utils.FileParserUtil;
import com.share.ai.common.utils.FileUtil;
import com.share.ai.common.utils.StrUtil;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/rga")
public class RgaController {
    // 支持的文件格式
    private static List<String> fileFormats = Arrays.asList(".txt", ".log", ".pdf", ".doc", ".docx", ".md");

    //向量数据库
    @Autowired
    public EmbeddingStore embeddingStore;


    //向量模型  （也就是把汉字变成数字的变换） ---stater 自动装配
    @Autowired
    public QwenEmbeddingModel qwenEmbeddingModel;


    //把字符串添加到向量
    @GetMapping("/addStr")
    public String addStr(String text) {
        // 将字符串转换为TextSegment
        TextSegment textSegment = TextSegment.from(text);

        // 把字符串生成向量
        Embedding embedding = qwenEmbeddingModel.embed(text).content();

        // 存储到向量数据库
        embeddingStore.add(embedding, textSegment);

        return "";
    }

    //把文件添加到向量
    @GetMapping("/addFile")
    public Integer addFile(String filePath) {
        if ((StrUtil.isNotBlank(filePath) && new File(filePath).exists())) {
            log.error("路径不存在filePath={}", filePath);
            return -1;
        }

        //文档筛选
        List<File> files = FileUtil.listFile(filePath);
        List<Document> documents = files.stream()
                .filter(f -> f.isFile() && fileFormats.stream().anyMatch(t -> f.getPath().endsWith(t)))
                .map(f -> loadDocumentWithParser(f))
                .filter(e -> e != null)
                .collect(Collectors.toList());

        DocumentSplitter splitter = DocumentSplitters.recursive(500, 100);
        int totalChunks = 0;
        for (Document document : documents) {
            try {
                // 1. 文件分块
                List<TextSegment> segments = splitter.split(document);
                totalChunks += segments.size();

                // 2. 为每个块生成嵌入向量
                List<Embedding> embeddings = segments.stream()
                        .map(e -> qwenEmbeddingModel.embed(e).content())
                        .collect(Collectors.toList());

                // 3. 存储到向量数据库
                int size = segments.size();
                for (int i = 0; i < size; i++) {
                    embeddingStore.add(embeddings.get(i), segments.get(i));
                }
            } catch (Exception e) {
                log.error("文档向量化失败: {}", document.metadata(), e);
            }
        }

        return totalChunks;
    }

    /**
     * 加载单个文件并自动选择解析器
     */
    private static Document loadDocumentWithParser(File file) {
        try {
            String fileName = file.getName().toLowerCase();
            String fileType = fileName.substring(fileName.lastIndexOf("."));
            DocumentParser parser = FileParserUtil.getDocumentParser(fileType);//定义转换器

            if (parser == null) {
                System.out.println("未找到文件 " + file.getName() + " 的解析器，跳过处理");
                return null;
            }

            return FileSystemDocumentLoader.loadDocument(file.toPath(), parser);

        } catch (Exception e) {
            System.err.println("解析文件失败: " + file.getName() + " - " + e.getMessage());
            return null;
        }
    }

    public static void main(String[] args) {
        new RgaController().addFile("D:\\data\\logs");
    }
}
